from pydantic import BaseModel, Field


class DrillingParameters(BaseModel):
    """钻井参数模型"""
    well_name: str = Field(..., description="井号")
    well_depth: float = Field(..., gt=0, description="井深（m）")
    workover_fluid_density: float = Field(..., gt=0, description="修井液密度（kg/m³）")
    steel_density: float = Field(7850, gt=0, description="钢材密度（kg/m³）")
    elastic_modulus: float = Field(..., gt=0, description="弹性模量（N/m²）")
    friction_coefficient: float = Field(..., ge=0, description="摩擦系数")
    safety_factor: float = Field(..., gt=0, description="大钩载荷（kN）")
    tensile_safety_factor: float = Field(..., gt=0, description="井口压力（MPa）")
    pulling_force: float = Field(..., description="上提拉力（kN）")
    collar_inner_outer_diameter: float = Field(..., description="钻具伸长度（mm）")
    cross_sectional_area: float = Field(..., gt=0, description="被卡物体截面积（mm²）")
    slip_length: float = Field(..., gt=0, description="井筒内径（mm）")
    collar_length: float = Field(..., gt=0, description="修井液动切力（Pa）")
    collar: float = Field(..., gt=0, description="修井液塑性粘度（MPa·s）")
    length: float = Field(..., gt=0, description="泵排量（L/min）")

    def to_si_units(self) -> dict:
        """转换为国际单位制"""
        return {
            "well_name": self.well_name,
            "well_depth": self.well_depth,  # m 已经是SI单位
            "workover_fluid_density": self.workover_fluid_density,  # kg/m³ 已经是SI单位
            "steel_density": self.steel_density,  # kg/m³ 已经是SI单位
            "elastic_modulus": self.elastic_modulus,  # N/m² 已经是SI单位
            "friction_coefficient": self.friction_coefficient,  # 无量纲
            "safety_factor": self.safety_factor * 1000,  # kN → N
            "tensile_safety_factor": self.tensile_safety_factor * 1e6,  # MPa → Pa
            "pulling_force": self.pulling_force * 1000,  # kN → N
            "collar_inner_outer_diameter": self.collar_inner_outer_diameter,  # 字符串保持不变
            "cross_sectional_area": self.cross_sectional_area * 1e-6,  # mm² → m²
            "slip_length": self.slip_length * 0.001,  # mm → m
            "collar_length": self.collar_length,  # Pa 已经是SI单位
            "collar": self.collar * 1e6,  # MPa·s → Pa·s
            "length": self.length / 60000,  # L/min → m³/s (1L=0.001m³, 1min=60s)
        }

    def to_dict(self) -> dict:
        """转换为字典格式"""
        return {
            "well_name": self.well_name,
            "well_depth": self.well_depth,
            "workover_fluid_density": self.workover_fluid_density,
            "steel_density": self.steel_density,
            "elastic_modulus": self.elastic_modulus,
            "friction_coefficient": self.friction_coefficient,
            "safety_factor": self.safety_factor,
            "tensile_safety_factor": self.tensile_safety_factor,
            "pulling_force": self.pulling_force,
            "collar_inner_outer_diameter": self.collar_inner_outer_diameter,
            "cross_sectional_area": self.cross_sectional_area,
            "slip_length": self.slip_length,
            "collar_length": self.collar_length,
            "collar": self.collar,
            "length": self.length,
        }

